package com.yc.cloud.openai.portal.service.impl;

import java.nio.file.Paths;
import java.time.LocalDateTime;
import java.util.Map;

import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import com.yc.cloud.common.basic.exception.Asserts;
import com.yc.cloud.common.basic.utils.PathUtils;
import com.yc.cloud.common.oss.service.OssFileService;
import com.yc.cloud.openai.entity.ImageRecord;
import com.yc.cloud.openai.portal.dao.ComfyUiClient;
import com.yc.cloud.openai.portal.dto.HistoryDto;
import com.yc.cloud.openai.portal.dto.ImageRecordNoteDto;
import com.yc.cloud.openai.portal.dto.PromptDto;
import com.yc.cloud.openai.portal.dto.QueueDto;
import com.yc.cloud.openai.portal.dto.UploadImageDto;
import com.yc.cloud.openai.portal.dto.request.TextToImageHistoryRequest;
import com.yc.cloud.openai.portal.dto.request.TextToImageInterruptRequest;
import com.yc.cloud.openai.portal.dto.request.TextToImageQueueCancelRequest;
import com.yc.cloud.openai.portal.dto.request.TextToImageQueueRequest;
import com.yc.cloud.openai.portal.dto.request.TextToImageRequest;
import com.yc.cloud.openai.portal.dto.request.TextToImageUploadRequest;
import com.yc.cloud.openai.portal.dto.response.ComfyUiQueueResponse;
import com.yc.cloud.openai.portal.enums.ImagePromptType;
import com.yc.cloud.openai.portal.enums.ImageRecordStatusEnums;
import com.yc.cloud.openai.portal.enums.PromptType;
import com.yc.cloud.openai.portal.enums.TextToImageProvider;
import com.yc.cloud.openai.portal.factory.TextToImageClientFactory;
import com.yc.cloud.openai.portal.service.ImageRecordService;
import com.yc.cloud.openai.portal.service.TextToImageService;

import cn.hutool.core.io.FileUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import lombok.SneakyThrows;
import lombok.val;
import lombok.extern.slf4j.Slf4j;

@Service
@Slf4j
public class TextToImageServiceImpl implements TextToImageService {

    // private final Map<String, TextToSpeechClient> ttsClients;
    private final Map<TextToImageProvider, ComfyUiClient> textToImageClientMap;
    private final ImageRecordService imageRecordService;
    private final OssFileService ossFileService;
    private final PathUtils pathUtils;
    public TextToImageServiceImpl(TextToImageClientFactory clientFactory, ImageRecordService imageRecordService,OssFileService ossFileService,PathUtils pathUtils) {
        this.imageRecordService = imageRecordService;
        this.ossFileService = ossFileService;
        this.pathUtils = pathUtils;
        // 使用枚举的 value 作为键来初始化 Map
        this.textToImageClientMap = clientFactory.getClientMap();
    }

    @Override
    public PromptDto convertTextToImage(TextToImageRequest request) {
        // 根据请求中的供应商选择对应的客户端
        val client = textToImageClientMap.get(request.getProvider());
        if (client == null) {
            Asserts.fail("Unsupported provider: " + request.getProvider());
            return null;
        }
        val textToImageResponse = client.convertTextToImage(request);
        if (textToImageResponse != null && StrUtil.isNotEmpty(textToImageResponse.getPromptId())) {
            // 保存文生图的一些信息到数据库中,便于后续查询
            val saveFlag = generateImageRecord(request, textToImageResponse);
            if (!saveFlag) {
                Asserts.fail("文生图记录保存失败");
            }
        }
        return textToImageResponse;
    }

    @Override
    public QueueDto queue(TextToImageQueueRequest request) {
        val client = getClient(request.getProvider());
        assert client != null;
        return client.queue(request);
    }

    /**
     * 获取指定任务ID的生图的历史记录
     *
     * @param request 请求参数
     * @return 历史记录
     */
    @Override
    public HistoryDto history(TextToImageHistoryRequest request) {
        val client = getClient(request.getProvider());
        assert client != null;
        val result = client.history(request.getPromptId(), PromptType.TEXT2IMAGE.getKey());
        // 将结果转换为响应对象
        val historyDto = new HistoryDto();
        historyDto.setHistoryList(result.getHistoryList());
        return historyDto;
    }

    @Override
    public ComfyUiQueueResponse cancelQueue(TextToImageQueueCancelRequest request) {
        val client = getClient(request.getProvider());
        assert client != null;
        return client.cancelQueue(request);
    }

    @Override
    public boolean interrupt(TextToImageInterruptRequest request) {
        val promptId = request.getPromptId();
        val client = getClient(request.getProvider());
        assert client != null;
        boolean result = false;
        // 先查看任务ID是否正在运行
        val isRunning = client.isRunningQueue(promptId);
        if (isRunning) {
            // 取消任务
            result = client.interrupt(request);
            if (result) {
                // 更新数据库中任务状态
                imageRecordService.updateStatusByUuid(promptId, ImageRecordStatusEnums.CANCELLED.getKey());
            }
        }
        log.error("interrupt failed, promptId:{}, not running.", promptId);
        return result;
    }

    private ComfyUiClient getClient(TextToImageProvider provider) {
        val client = textToImageClientMap.get(provider);
        if (client == null) {
            Asserts.fail("Unsupported provider: " + provider.getValue());
            return null;
        }
        return client;
    }

    // 生成文生图的记录
    private boolean generateImageRecord(TextToImageRequest textToImageRequest, PromptDto promptDto) {
        val imageRecord = new ImageRecord();
        // 客户端ID
        imageRecord.setClientId(promptDto.getClientId());
        // 提示词内容
        imageRecord.setPrompt(textToImageRequest.getPrompt());
        // 设置任务ID
        imageRecord.setUuid(promptDto.getPromptId());
        // 设置生成图片的类型为文生图枚举值
        imageRecord.setType(ImagePromptType.TEXT2IMAGE.getKey());
        imageRecord.setCreateTime(LocalDateTime.now());
        imageRecord.setProvider(textToImageRequest.getProvider().getKey());

        // 设置一些额外信息,如图片尺寸、生图模式等信息
        val imageRecordNoteDto = new ImageRecordNoteDto();
        imageRecordNoteDto.setWidth(textToImageRequest.getWidth());
        imageRecordNoteDto.setHeight(textToImageRequest.getHeight());
        imageRecordNoteDto.setMode(textToImageRequest.getMode());
        imageRecordNoteDto.setProvider(textToImageRequest.getProvider().getValue());
        imageRecord.setNote(JSONUtil.toJsonStr(imageRecordNoteDto));
        return imageRecordService.save(imageRecord);
    }

    @Override
    @SneakyThrows
    public UploadImageDto upload(TextToImageUploadRequest request, MultipartFile multipartFile) {
        val client = getClient(request.getProvider());
        assert client != null;
        
        val path = Paths.get(pathUtils.getBasePath(), multipartFile.getOriginalFilename());
        // 创建临时文件
        val tempFile = FileUtil.writeBytes(multipartFile.getBytes(), path.toFile());
        return client.uploadImage(tempFile);
    }

}
